#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-i', '--input', help='Input file', default='mz.180x90x30.h0.nc')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-l', '--loop', action='store_true', help='Loop the plot')
parser.add_argument('-u', '--update-seconds', type=float, default=10, help='Update seconds')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

fig = plt.figure(figsize=(10, 6))
proj = ccrs.PlateCarree(central_longitude=180)

def plot(data_path, fig):
  f = addfile(data_path)

  p = f.ph[args.time_step,:,:,:]

  fig.suptitle(((f.time[args.time_step] - f.time[0]).values).astype('timedelta64[D]'), fontsize=16)

  i = 0

  i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
  u700 = vinterp_p(p, f.u[args.time_step,:,:,:], 700 * units.hPa)
  plot_contour_map(ax, u700[0,:,:], cmap='jet', levels=np.linspace(-15, 50, 14), left_string='U at 700 hPa')

  i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
  v700 = vinterp_p(p, f.v[args.time_step,:,:,:], 700 * units.hPa)
  plot_contour_map(ax, v700[0,:,:], cmap='jet', left_string='V at 700 hPa', levels=np.linspace(-35, 20, 12))

  i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
  z700 = vinterp_p(p, f.gz[args.time_step,:,:,:], 700 * units.hPa) / 9.80616
  plot_contour_map(ax, z700[0,:,:], cmap='jet', left_string='Z at 700 hPa', levels=np.linspace(2400, 3400, 11))

  i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
  t700 = vinterp_p(p, f.t[args.time_step,:,:,:], 700 * units.hPa)
  plot_contour_map(ax, t700[0,:,:], cmap='jet', left_string='T at 700 hPa', levels=np.linspace(270, 303, 12))

plot(args.input, fig)
if args.loop:
  loop_plot(args.input, fig, plot, args.update_seconds)
elif args.output:
  plt.savefig(args.output, dpi=300)
else:
  plt.tight_layout()
  plt.show()
